"""
TCP连接管理服务
"""
import socket
import select
from typing import Dict, Optional
from models import TCPConfig, ConnectionStats, APIResponse

class TCPConnectionManager:
    """TCP连接管理器"""
    
    def __init__(self):
        self.connections: Dict[str, dict] = {}
    
    def create_connection(self, config: TCPConfig) -> APIResponse:
        """
        创建TCP连接
        
        Args:
            config: TCP配置
            
        Returns:
            APIResponse: 连接结果
        """
        connection_id = f"{config.host}:{config.port}"
        
        try:
            if config.mode == "client":
                sock = self._create_client_connection(config)
            else:  # server mode
                sock = self._create_server_connection(config)
                
            self.connections[connection_id] = {
                "socket": sock,
                "config": config.dict(),
                "connected": True,
                "stats": ConnectionStats().dict(),
                "clients": [] if config.mode == "server" else None
            }
            
            message = f"成功连接到 {config.host}:{config.port}" if config.mode == "client" else f"服务器已启动，监听 {config.host}:{config.port}"
            
            # 返回的响应应该同时在顶级和data中包含connection_id
            response_data = {
                "success": True,
                "message": message,
                "data": {"connection_id": connection_id},
                "connection_id": connection_id  # 添加顶级字段兼容性
            }
            
            return APIResponse(**response_data)
            
        except Exception as e:
            return APIResponse(
                success=False,
                message=f"连接失败: {str(e)}"
            )
    
    def _create_client_connection(self, config: TCPConfig) -> socket.socket:
        """创建客户端连接"""
        if config.protocol == "tcp":
            sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            sock.connect((config.host, config.port))
        else:  # UDP
            sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        return sock
    
    def _create_server_connection(self, config: TCPConfig) -> socket.socket:
        """创建服务器连接"""
        sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        sock.bind((config.host, config.port))
        sock.listen(5)
        # 设置非阻塞模式用于accept
        sock.settimeout(0.1)
        return sock
    
    def disconnect(self, connection_id: str) -> APIResponse:
        """
        断开连接
        
        Args:
            connection_id: 连接ID
            
        Returns:
            APIResponse: 断开结果
        """
        if connection_id not in self.connections:
            return APIResponse(
                success=False,
                message="连接不存在"
            )
        
        try:
            conn = self.connections[connection_id]
            
            # 关闭客户端连接（如果是服务器模式）
            if conn.get("clients"):
                for client in conn["clients"]:
                    client.close()
            
            # 关闭主socket
            conn["socket"].close()
            del self.connections[connection_id]
            
            return APIResponse(
                success=True,
                message="连接已断开"
            )
            
        except Exception as e:
            return APIResponse(
                success=False,
                message=f"断开连接失败: {str(e)}"
            )
    
    def send_data(self, connection_id: str, data: bytes) -> APIResponse:
        """
        发送数据
        
        Args:
            connection_id: 连接ID
            data: 要发送的数据
            
        Returns:
            APIResponse: 发送结果
        """
        if connection_id not in self.connections:
            return APIResponse(
                success=False,
                message="连接不存在"
            )
        
        try:
            conn = self.connections[connection_id]
            
            if conn["config"]["mode"] == "client":
                conn["socket"].send(data)
            else:
                # 服务器模式，发送给所有客户端
                for client in conn.get("clients", []):
                    client.send(data)
            
            # 更新统计信息
            conn["stats"]["send_count"] += 1
            conn["stats"]["send_bytes"] += len(data)
            
            return APIResponse(
                success=True,
                message="数据发送成功"
            )
            
        except Exception as e:
            return APIResponse(
                success=False,
                message=f"发送失败: {str(e)}"
            )
    
    def read_data(self, connection_id: str) -> Optional[bytes]:
        """
        读取TCP数据
        
        Args:
            connection_id: 连接ID
            
        Returns:
            读取到的数据或None
        """
        if connection_id not in self.connections:
            return None
        
        try:
            conn = self.connections[connection_id]
            sock = conn["socket"]
            
            # 设置非阻塞模式
            sock.settimeout(0.1)  # 100ms超时
            
            if conn["config"]["mode"] == "client":
                # 客户端模式：从服务器接收数据
                try:
                    data = sock.recv(4096)
                    if data:
                        # 更新统计信息
                        conn["stats"]["receive_bytes"] += len(data)
                        conn["stats"]["receive_count"] += 1
                        return data
                except socket.timeout:
                    pass
                except Exception:
                    pass
            else:
                # 服务器模式：检查是否有新客户端连接
                try:
                    client_sock, addr = sock.accept()
                    client_sock.settimeout(0.1)
                    conn["clients"].append(client_sock)
                    print(f"新客户端连接: {addr}")
                    
                    # 尝试从新客户端读取数据
                    try:
                        data = client_sock.recv(4096)
                        if data:
                            conn["stats"]["receive_bytes"] += len(data)
                            conn["stats"]["receive_count"] += 1
                            return data
                    except socket.timeout:
                        pass
                except socket.timeout:
                    # 没有新连接，检查现有客户端
                    for i, client in enumerate(conn.get("clients", [])):
                        try:
                            data = client.recv(4096)
                            if data:
                                conn["stats"]["receive_bytes"] += len(data)
                                conn["stats"]["receive_count"] += 1
                                return data
                        except socket.timeout:
                            continue
                        except Exception:
                            # 客户端断开，移除
                            try:
                                client.close()
                            except:
                                pass
                            conn["clients"].pop(i)
                            break
                except Exception:
                    pass
            
        except Exception as e:
            # 更新错误统计
            if connection_id in self.connections:
                self.connections[connection_id]["stats"]["error_count"] += 1
        
        return None

    def get_connection_status(self, connection_id: str) -> Optional[dict]:
        """
        获取连接状态
        
        Args:
            connection_id: 连接ID
            
        Returns:
            连接状态信息或None
        """
        if connection_id in self.connections:
            conn = self.connections[connection_id]
            return {
                "connected": conn["connected"],
                "config": conn["config"],
                "stats": conn["stats"]
            }
        return None
    
    def get_connection(self, connection_id: str) -> Optional[dict]:
        """获取连接对象"""
        return self.connections.get(connection_id)

# 全局TCP连接管理器实例
tcp_manager = TCPConnectionManager()